#!/usr/bin/env python
# -*- coding:utf-8 -*-
from controller.apiKeyController import ApiKeyController
from controller.ssoUserController import SsoUserController
from llm.model import LLMFactory
from utils.response_code import RET
from utils.rsa_encryption_decryption import RSAEncryptionDecryption


class LLMService:

    @classmethod
    def get_model(cls):
        """
            获取LLM的chat 方法
        """
        from flask import g
        user = SsoUserController.get(user_id=g.user['user_id'])
        if user['code'] != RET.OK:
            raise Exception('user not found')
        if len(user['data']) != 1:
            raise Exception('user not found')
        user = user['data'][0]

        if user['key_id'] is None:
            raise Exception('API KEY not set')

        api_key = ApiKeyController.get(key_id=user['key_id'])
        if api_key['code'] != RET.OK:
            raise Exception('API KEY not found')
        if len(api_key['data']) != 1:
            raise Exception('API KEY not found')
        api_key = api_key['data'][0]
        model = LLMFactory.get_model(model_name=api_key['chat_model'],
                                     api_key=RSAEncryptionDecryption.decrypt(api_key['content']))
        return model
